#pragma once

#include <iostream>

using namespace std;

enum Color
{
    RED, BLACK
};

template<class T>
struct RBTreeNode
{
    explicit RBTreeNode (const T &data = T (), Color color = RED)
            : pLeft (nullptr), pRight (nullptr), parent (nullptr), data (data), color (color)
    {}
    
    RBTreeNode<T> *pLeft;
    RBTreeNode<T> *pRight;
    RBTreeNode<T> *parent;
    T             data;
    Color         color;
};

template<class T>
class RBTree
{
    typedef RBTreeNode<T> Node;
public:
    RBTree ()
    {
        pHead = new Node;
        pHead->pLeft  = pHead;
        pHead->pRight = pHead;
    }
    
    bool Insert (const T &data)
    {
        // 按照二叉搜索树的性质插入新节点
        Node *&pRoot = GetRoot ();
        if (nullptr == pRoot)
        {
            pRoot = new Node (data, BLACK);
            pRoot->parent = pHead;
            return true;
        }
        else
        {
            // 找待插入节点在二叉搜索树中的位置
            // 并保存其双亲节点
            Node *pCur    = pRoot;
            Node *pParent = nullptr;
            while (pCur)
            {
                pParent = pCur;
                if (data < pCur->data)
                {
                    pCur = pCur->pLeft;
                }
                else if (data > pCur->data)
                {
                    pCur = pCur->pRight;
                }
                else
                {
                    return false;
                }
            }
            
            // 插入新节点
            pCur = new Node (data);
            if (data < pParent->data)
            {
                pParent->pLeft = pCur;
            }
            else
            {
                pParent->pRight = pCur;
            }
            
            pCur->parent = pParent;
            
            // pCur新节点默认颜色：红色
            // 如果pParent的颜色是黑色的，满足红黑树性质
            // 如果pParent的颜色是红色的，违反了性质三--不能有连在一起的红色节点
            while (pParent != pHead && pParent->color == RED)
            {
                Node *grandFather = pParent->parent;
                if (pParent == grandFather->pLeft)
                {
                    Node *uncle = grandFather->pRight;
                    // 叔叔节点存在且为红
                    if (uncle && RED == uncle->color)
                    {
                        // 情况一
                        pParent->color     = BLACK;
                        uncle->color       = BLACK;
                        grandFather->color = RED;
                        pCur    = grandFather;
                        pParent = pCur->parent;
                    }
                    else
                    {
                        // 叔叔节点不存在 || 叔叔节点存在 && 黑色
                        if (pCur == pParent->pRight)
                        {
                            // 情况三
                            RotateL (pParent);
                            swap (pCur, pParent);
                        }
                        
                        // 情况二：
                        pParent->color     = BLACK;
                        grandFather->color = RED;
                        RotateR (grandFather);
                    }
                }
                else
                {
                    // 一二三的反情况
                    Node *uncle = grandFather->pLeft;
                    if (uncle && RED == uncle->color)
                    {
                        pParent->color     = BLACK;
                        uncle->color       = BLACK;
                        grandFather->color = RED;
                        pCur    = grandFather;
                        pParent = pCur->parent;
                    }
                    else
                    {
                        // 叔叔节点不存在 || 存在且为黑
                        if (pCur == pParent->pLeft)
                        {
                            // 情况三反情况
                            RotateR (pParent);
                            swap (pCur, pParent);
                        }
                        
                        // 情况二的反情况
                        pParent->color     = BLACK;
                        grandFather->color = RED;
                        RotateL (grandFather);
                    }
                }
            }
        }
        
        pRoot->color  = BLACK;
        pHead->pLeft  = LeftMost ();
        pHead->pRight = RightMost ();
        return true;
    }
    
    Node *&GetRoot ()
    {
        return pHead->parent;
    }
    
    void InOrder ()
    {
        InOrder (GetRoot ());
    }
    
    Node *LeftMost ()
    {
        Node *pRoot = GetRoot ();
        if (nullptr == pRoot)
        {
            return pHead;
        }
        
        Node *pCur = pRoot;
        while (pCur->pLeft)
        {
            pCur = pCur->pLeft;
        }
        
        return pCur;
    }
    
    Node *RightMost ()
    {
        Node *pRoot = GetRoot ();
        if (nullptr == pRoot)
        {
            return pHead;
        }
        
        Node *pCur = pRoot;
        while (pCur->pRight)
        {
            pCur = pCur->pRight;
        }
        
        return pCur;
    }
    
    bool IsValidRBTRee ()
    {
        Node *pRoot = GetRoot ();
        if (nullptr == pRoot)
        {
            return true;
        }
        
        if (pRoot->color != BLACK)
        {
            cout << "违反性质一：根节点是黑色" << endl;
            return false;
        }
        
        // 获取一条路径中黑色节点的个数
        size_t blackCount = 0;
        Node   *pCur      = pRoot;
        while (pCur)
        {
            if (pCur->color == BLACK)
            {
                blackCount++;
            }
            
            pCur = pCur->pLeft;
        }
        
        size_t pathBlack = 0;
        return IsValidRBTRee (pRoot, blackCount, pathBlack);
    }

private:
    bool IsValidRBTRee (Node *pRoot, size_t blackCount, size_t pathBlack)
    {
        if (nullptr == pRoot)
        {
            return true;
        }
        
        if (pRoot->color == BLACK)
        {
            pathBlack++;
        }
        
        Node *pParent = pRoot->parent;
        if (pParent != pHead &&
            RED == pParent->color && RED == pRoot->color)
        {
            cout << "违反性质3：不能存在连在一起的红色节点" << endl;
            return false;
        }
        
        // 叶子节点
        if (nullptr == pRoot->pLeft && nullptr == pRoot->pRight)
        {
            if (pathBlack != blackCount)
            {
                cout << "违反性质4：每条路径中黑色节点个数均相同" << endl;
                return false;
            }
        }
        
        return IsValidRBTRee (pRoot->pLeft, blackCount, pathBlack) &&
               IsValidRBTRee (pRoot->pRight, blackCount, pathBlack);
    }
    
    void InOrder (Node *pRoot)
    {
        if (pRoot)
        {
            InOrder (pRoot->pLeft);
            cout << pRoot->data << " ";
            InOrder (pRoot->pRight);
        }
    }
    
    void RotateL (Node *pParent)
    {
        Node *pSubR  = pParent->pRight;
        Node *pSubRL = pSubR->pLeft;
        
        pParent->pRight = pSubRL;
        if (pSubRL)
        {
            pSubRL->parent = pParent;
        }
        
        pSubR->pLeft = pParent;
        Node *pPParent = pParent->parent;
        pSubR->parent   = pPParent;
        pParent->parent = pSubR;
        
        if (pPParent == pHead)
        {
            pHead->parent = pSubR;
        }
        else
        {
            if (pParent == pPParent->pLeft)
            {
                pPParent->pLeft = pSubR;
            }
            else
            {
                pPParent->pRight = pSubR;
            }
        }
    }
    
    void RotateR (Node *pParent)
    {
        Node *pSubL  = pParent->pLeft;
        Node *pSubLR = pSubL->pRight;
        
        pParent->pLeft = pSubLR;
        if (pSubLR)
        {
            pSubLR->parent = pParent;
        }
        
        pSubL->pRight = pParent;
        
        Node *pPParent = pParent->parent;
        pParent->parent = pSubL;
        pSubL->parent   = pPParent;
        
        if (pPParent == pHead)
        {
            pHead->parent = pSubL;
        }
        else
        {
            if (pParent == pPParent->pLeft)
            {
                pPParent->pLeft = pSubL;
            }
            else
            {
                pPParent->pRight = pSubL;
            }
        }
    }

private:
    Node *pHead;
};
